##---------------------------------------------------------------
## Required libraries for computing
##---------------------------------------------------------------

#----- Parallel computing
library(doParallel)

#----- Make clusters based on the number of CPU cores
cl <- makeCluster(detectCores())
registerDoParallel(cl)
getDoParWorkers()

#----- Support parallel excution
library(foreach)

#----- Stan (HMC) for R
library(rstan)

##---------------------------------------------------------------
## Extract MCMC samples from Stan outputs
##---------------------------------------------------------------

#----- Load MCMC samples and Data
load("MCMCsample.RData")
load("Master.RData")

#------ Set Treatment (TRT), Outcome (OUT), Mediators (M) and Covariates (X)
Data <- Master
OUT <- Data$PM.2.5
TRT <- Data$SO2.SC
M <- cbind(Data$SO2_Annual, Data$NOx_Annual, Data$CO2_Annual)
X <- cbind(Data$S_n_CR, Data$NumNOxControls, Data$Heat_Input/100000, Data$Barometric_Pressure, Data$Temperature,  Data$PctCapacity, Data$sulfur_Content, Data$Phase2_Indicator, Data$Operating_Time/1000)

dim.cov <- dim(X)[2] #<--------- Num. of Covariates

#------ Variables by treatments
x0 <- X[which(TRT==0),]
x1 <- X[which(TRT==1),]

y0 <- OUT[which(TRT==0)]
y1 <- OUT[which(TRT==1)]

m0 <- log(M[which(TRT==0),])
m1 <- log(M[which(TRT==1),])

n0 <- dim(x0)[1]
n1 <- dim(x1)[1]



#----- Extract MCMC samples from each marginal distribution

data.y0 <- extract(fit.y0)
data.y1 <- extract(fit.y1)
data.m10 <- extract(fit.m1_0)
data.m11 <- extract(fit.m1_1)
data.m20 <- extract(fit.m2_0)
data.m21 <- extract(fit.m2_1)
data.m30 <- extract(fit.m3_0)
data.m31 <- extract(fit.m3_1)

#----- Parameters of the models under Z=1

gamma10 <- data.y1$gamma0
gamma11 <- data.y1$gamma1
w1 <- data.y1$W
beta1_10 <- data.m11$gamma0
beta1_11 <- data.m11$gamma1
beta2_10 <- data.m21$gamma0
beta2_11 <- data.m21$gamma1
beta3_10 <- data.m31$gamma0
beta3_11 <- data.m31$gamma1
s1_1 <- data.m11$W
s2_1 <- data.m21$W
s3_1 <- data.m31$W
psi1 <- data.y1$psi
psi1_1 <- data.m11$psi
psi2_1 <- data.m21$psi
psi3_1 <- data.m31$psi


#----- Parameters of the models under Z=0

gamma00 <- data.y0$gamma0
gamma01 <- data.y0$gamma1
w0 <- data.y0$W
beta1_00 <- data.m10$gamma0
beta1_01 <- data.m10$gamma1
beta2_00 <- data.m20$gamma0
beta2_01 <- data.m20$gamma1
beta3_00 <- data.m30$gamma0
beta3_01 <- data.m30$gamma1
s1_0 <- data.m10$W
s2_0 <- data.m20$W
s3_0 <- data.m30$W
psi0 <- data.y0$psi
psi1_0 <- data.m10$psi
psi2_0 <- data.m20$psi
psi3_0 <- data.m30$psi


#----- Setting Thinning and Burn-in's
Thin <- 5
Burn0 <- 18000  # Extra burn-in periods for Models under Z=0
Burn1 <- 18000   # Extra burn-in periods for Models under Z=1

#----- Set the number of post-processing steps, N
n.iter <- 400


##---------------------------------------------------------------
## 'Main' function on each cluster (parallel)
##---------------------------------------------------------------


main<-function(temp){


    #----- Require to fit multivariate normal distributions
    library(mnormt)
    joint0 <- cbind(y0, m0);  joint1 <- cbind(y1, m1)
    #----- Estimating condtional correlations
    fit0 <- lm(joint0 ~ x0)
    fit1 <- lm(joint1 ~ x1)
    
    cov0 <- 1 / n0 * t(joint0 - predict(fit0, newdata=data.frame(x0))) %*% (joint0-predict(fit0, newdata = data.frame(x0)))
    cov1 <- 1 / n1 * t(joint1 - predict(fit1, newdata=data.frame(x1))) %*% (joint1-predict(fit1, newdata = data.frame(x1)))
    
    #----- Under Normality, the following relatioship between
    #----- Spearman and Pearson correlations holds (Kruskal, 1958)
    cor0 <- 6 / pi * asin(matrix(apply(expand.grid(1:4,1:4), 1, function(x) cov0[x[1],x[2]] / sqrt(cov0[x[1],x[1]] * cov0[x[2],x[2]])),4,4)/2)
    cor1 <- 6 / pi * asin(matrix(apply(expand.grid(1:4,1:4), 1, function(x) cov1[x[1],x[2]] / sqrt(cov1[x[1],x[1]]*cov1[x[2],x[2]])),4,4)/2)
    
    
    #----------------- Sensitivity parameters
    sens1 <- function(j,k) 2 * min(abs(cor0[j+1,k+1]), abs(cor1[j+1,k+1])) / abs(cor0[j+1,k+1] + cor1[j+1,k+1])
    rm <- function(j,k,rho1) (cor0[j+1,k+1] + cor1[j+1,k+1]) / 2 * rho1
    
    sens2 <- function(k) 2 * min(abs(cor0[1,k+1]), abs(cor1[1,k+1])) / abs(cor0[1,k+1]+ cor1[1,k+1])
    ry <- function(k,rho2) (cor0[1,k+1]+ cor1[1,k+1]) / 2 * rho2
    rho1 <- runif(1, 0, min(apply(expand.grid(1:3,1:3),1, function(x) sens1(x[1],x[2]))))
    rho2 <- runif(1, 0, min(sens2(1),sens2(2), sens2(3)))
    
    #----- Construct the correlation matrix
    cor.part <- matrix(apply(expand.grid(1:3,1:3), 1, function(x) rm(x[1], x[2], rho1)), 3, 3)
    cor.y1m0 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
    cor.y0m1 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
    
    COR1 <- 2 * sin(pi * (1/6) *rbind(cbind(cor1,rbind(cor.y1m0, cor.part)), cbind(t(rbind(cor.y1m0, cor.part)), cor0[2:4,2:4] ) ))
    COR0 <- 2 * sin(pi * (1/6) *rbind(cbind(cor0,rbind(cor.y0m1, cor.part)), cbind(t(rbind(cor.y0m1, cor.part)), cor1[2:4,2:4] ) ))
    
    #----- Check positive definiteness
    S0 <- attr(chol(COR0, pivot=TRUE),"rank");     S1 <- attr(chol(COR1, pivot=TRUE),"rank")
    while (S0 !=7 | S1 !=7){
        rho1 <- runif(1, 0, min(apply(expand.grid(1:3,1:3),1, function(x) sens1(x[1],x[2]))))
        rho2 <- runif(1, 0, min(sens2(1),sens2(2), sens2(3)))
        cor.part <- matrix(apply(expand.grid(1:3,1:3), 1, function(x) rm(x[1], x[2], rho1)), 3, 3)
        cor.y1m0 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
        cor.y0m1 <- matrix(sapply(1:3, function(x) ry(x,rho2)), nrow=1)
        COR1 <- 2 * sin(pi * (1/6) *rbind(cbind(cor1,rbind(cor.y1m0, cor.part)), cbind(t(rbind(cor.y1m0, cor.part)), cor0[2:4,2:4] ) ))
        COR0 <- 2 * sin(pi * (1/6) *rbind(cbind(cor0,rbind(cor.y0m1, cor.part)), cbind(t(rbind(cor.y0m1, cor.part)), cor1[2:4,2:4] ) ))
        S0 <- attr(chol(COR0, pivot=TRUE),"rank");     S1 <- attr(chol(COR1, pivot=TRUE),"rank")
    }
    
    #----- Index of iterations (on parallel)
    j <- temp
    
    #----- Index of posterior samples
    index0 <- Thin * j + Burn0; index1 <- Thin * j + Burn1
    
    #----- Empirical distributions for covariates
    size <- 15000
    s.covariate <- data.frame(X[sample(seq(1,dim(X)[1]), size = size, replace = TRUE),])  #----- Covariates samples
    
    
    #----- Weighting parameters of mixtures
    pi.y0 <- pmax(psi0[index0,], 0)
    pi.y1 <- pmax(psi1[index1,], 0)
    pi.m10 <- pmax(psi1_0[index0,], 0)
    pi.m20 <- pmax(psi2_0[index0,], 0)
    pi.m30 <- pmax(psi3_0[index0,], 0)
    pi.m11 <- pmax(psi1_1[index1,], 0)
    pi.m21 <- pmax(psi2_1[index1,], 0)
    pi.m31 <- pmax(psi3_1[index1,], 0)
    

    #----- Sampling Distribution for Y(1;M(0,0,0),M(1,1,1))
    f1.M0M1 <- function(m1, m2, m3, H){
        Cor01 <- COR1
        s.m1_0 <- qnorm(sum(apply(cbind(1:10,pi.m10), 1, function(x) x[2]*pnorm(m1, mean=beta1_00[index0,x[1]]+beta1_01[index0,]%*%H, sd=1/sqrt(s1_0[index0,x[1]])))),0,1)
        s.m2_0 <- qnorm(sum(apply(cbind(1:10,pi.m20), 1, function(x) x[2]*pnorm(m2, mean=beta2_00[index0,x[1]]+beta2_01[index0,]%*%H, sd=1/sqrt(s2_0[index0,x[1]])))),0,1)
        s.m3_0 <- qnorm(sum(apply(cbind(1:10,pi.m30), 1, function(x) x[2]*pnorm(m3, mean=beta3_00[index0,x[1]]+beta3_01[index0,]%*%H, sd=1/sqrt(s3_0[index0,x[1]])))),0,1)

        F <- pnorm(rmnorm(1, mean=Cor01[1:4,5:7]%*%solve(Cor01[5:7,5:7])%*%c(s.m1_0,s.m2_0,s.m3_0), Cor01[1:4,1:4]-Cor01[1:4,5:7]%*%solve(Cor01[5:7,5:7])%*%Cor01[5:7,1:4]),0,1)

        Fy1 <- function(hh,u){abs(sum(apply(cbind(1:10,pi.y1), 1, function(x) x[2]*pnorm(hh, mean=gamma10[index1,x[1]]+gamma11[index1,]%*%H, sd=1/sqrt(w1[index1,x[1]]))))-u)}
        Fm11 <- function(hh,u){abs(sum(apply(cbind(1:10,pi.m11), 1, function(x) x[2]*pnorm(hh, mean=beta1_10[index1,x[1]]+beta1_11[index1,]%*%H, sd=1/sqrt(s1_1[index1,x[1]]))))-u)}
        Fm21 <- function(hh,u){abs(sum(apply(cbind(1:10,pi.m21), 1, function(x) x[2]*pnorm(hh, mean=beta2_10[index1,x[1]]+beta2_11[index1,]%*%H, sd=1/sqrt(s2_1[index1,x[1]]))))-u)}
        Fm31 <- function(hh,u){abs(sum(apply(cbind(1:10,pi.m31), 1, function(x) x[2]*pnorm(hh, mean=beta3_10[index1,x[1]]+beta3_11[index1,]%*%H, sd=1/sqrt(s3_1[index1,x[1]]))))-u)}

        s.y1 <- optimize(Fy1, u=F[1], interval=c(3, 20))$minimum
        s.m1_1 <- optimize(Fm11,  u=F[2], interval=c(0, 10))$minimum
        s.m2_1 <- optimize(Fm21,  u=F[3], interval=c(0, 10))$minimum
        s.m3_1 <- optimize(Fm31,  u=F[4], interval=c(5, 15))$minimum

        return(c(s.y1,s.m1_1,s.m2_1,s.m3_1))
    }



    #----- Sampling Distribution for Y(0;M(0,0,0),M(1,1,1))
    f0.M0M1 <- function(m1, m2, m3, H){
        Cor01 <- COR0
    
        s.m1_1 <- qnorm(sum(apply(cbind(1:10,pi.m11), 1, function(x) x[2]*pnorm(m1, mean=beta1_10[index1,x[1]]+beta1_11[index1,]%*%H, sd=1/sqrt(s1_1[index1,x[1]])))),0,1)
        s.m2_1 <- qnorm(sum(apply(cbind(1:10,pi.m21), 1, function(x) x[2]*pnorm(m2, mean=beta2_10[index1,x[1]]+beta2_11[index1,]%*%H, sd=1/sqrt(s2_1[index1,x[1]])))),0,1)
        s.m3_1 <- qnorm(sum(apply(cbind(1:10,pi.m31), 1, function(x) x[2]*pnorm(m3, mean=beta3_10[index1,x[1]]+beta3_11[index1,]%*%H, sd=1/sqrt(s3_1[index1,x[1]])))),0,1)
    
        F <- pnorm(rmnorm(1, mean=Cor01[1:4,5:7]%*%solve(Cor01[5:7,5:7])%*%c(s.m1_1,s.m2_1,s.m3_1), Cor01[1:4,1:4]-Cor01[1:4,5:7]%*%solve(Cor01[5:7,5:7])%*%Cor01[5:7,1:4]),0,1)
    
        Fy0 <- function(hh,u){abs(sum(apply(cbind(1:10,pi.y0), 1, function(x) x[2]*pnorm(hh, mean=gamma00[index0,x[1]]+gamma01[index0,]%*%H, sd=1/sqrt(w0[index0,x[1]]))))-u)}
        Fm10 <- function(hh,u){abs(sum(apply(cbind(1:10,pi.m10), 1, function(x) x[2]*pnorm(hh, mean=beta1_00[index0,x[1]]+beta1_01[index0,]%*%H, sd=1/sqrt(s1_0[index0,x[1]]))))-u)}
        Fm20 <- function(hh,u){abs(sum(apply(cbind(1:10,pi.m20), 1, function(x) x[2]*pnorm(hh, mean=beta2_00[index0,x[1]]+beta2_01[index0,]%*%H, sd=1/sqrt(s2_0[index0,x[1]]))))-u)}
        Fm30 <- function(hh,u){abs(sum(apply(cbind(1:10,pi.m30), 1, function(x) x[2]*pnorm(hh, mean=beta3_00[index0,x[1]]+beta3_01[index0,]%*%H, sd=1/sqrt(s3_0[index0,x[1]]))))-u)}
    
        s.y0 <- optimize(Fy0, u=F[1], interval=c(3, 20))$minimum
        s.m1_0 <- optimize(Fm10,  u=F[2], interval=c(0, 10))$minimum
        s.m2_0 <- optimize(Fm20,  u=F[3], interval=c(0, 10))$minimum
        s.m3_0 <- optimize(Fm30,  u=F[4], interval=c(5, 15))$minimum
        return(c(s.y0,s.m1_0,s.m2_0,s.m3_0))
    }


    #----- Impute Y samples
    new.data <- cbind(TRT,OUT,log(M),X)
    new.data1 <- array(dim=c(dim(new.data)[1],4,40))
    for(j in 1:40){
        new.data1[,,j] <- t(apply(new.data, 1, function(x) x[1]*f0.M0M1(x[3],x[4],x[5],x[6:13])+(1-x[1])*f1.M0M1(x[3],x[4],x[5],x[6:13])))
    }
    data.list <- list(new.data1[,,1],new.data1[,,2],new.data1[,,3],new.data1[,,4],new.data1[,,5],new.data1[,,6],new.data1[,,7],new.data1[,,8],new.data1[,,9],new.data1[,,10],new.data1[,,11],new.data1[,,12],new.data1[,,13],new.data1[,,14],new.data1[,,15],new.data1[,,16],new.data1[,,17],new.data1[,,18],new.data1[,,19],new.data1[,,20],new.data1[,,21],new.data1[,,22],new.data1[,,23],new.data1[,,24],new.data1[,,25],new.data1[,,26],new.data1[,,27],new.data1[,,28],new.data1[,,29],new.data1[,,30],new.data1[,,31],new.data1[,,32],new.data1[,,33],new.data1[,,34],new.data1[,,35],new.data1[,,36],new.data1[,,37],new.data1[,,38],new.data1[,,39],new.data1[,,40])
    new.data2<-cbind(new.data, Reduce("+", data.list) / length(data.list))

    #----- Complete data
    y1.sample <- new.data2[which(new.data2[,1]==1),c(2,14,15,16,17,3,4,5)]
    y0.sample <- new.data2[which(new.data2[,1]==0),c(14,2,3,4,5,15,16,17)]
    y.sample <- rbind(y1.sample, y0.sample)
    
 

    return(y.sample)
}

result<-foreach(temp = 1:n.iter, .combine = rbind) %dopar% main(temp)

save(result, file="Ysample.RData")
stopCluster(cl)


